# !usr/bin/env python
# -*- coding:utf-8 _*-

"""
Ref:
    Bilibili: https://www.bilibili.com/video/BV1u54y1Q7qf/?spm_id_from=333.337.search-card.all.click
    CSDN: https://blog.csdn.net/GrayOnDream/article/details/124564129
    CSDN: https://blog.csdn.net/qq_42001367/article/details/107958821
"""


from elftools.dwarf.descriptions import set_global_machine_arch
from elftools.elf.elffile import ELFFile
from dataclasses import dataclass, field
from Logger import LOG


@dataclass
class Item:
    name: str = ""  # 变量名
    addr: int = 0  # 变量地址
    data_type: str = "unsigned char"  # 数据类型
    data_type_size: int = 1  # 数据类型占用字节数，方便计算不同角标的偏移
    array_length: int = 0  # 数组长度，0表示非数组
    enum_table: dict = None  # 枚举类型，key为枚举值，value为枚举名称  {0：'OFF', 1：'ON'}
    member_offset: int = 0  # 成员偏移字节数，用于数组下标计算
    parent: "Item" = None  # 父节点
    children: list["Item"] = field(default_factory=list)  # 子节点
    level: int = 0  # 节点层级，0表示根节点
    attributes: list = field(default_factory=list)  # 属性列表，用于记录一些附加信息，如是否是静态变量等

    def append(self, child):
        self.children.append(child)
        # setattr(self, child.name, child)  # child.name中的名称会覆盖掉原有的属性，比如addr。

    def __str__(self):
        tab = "    "
        s = (tab * self.level +
             f"level: {self.level} "
             f"name: {self.name:<32} "
             f"addr: {hex(self.addr):<8} "
             f"data_type_size: {self.data_type_size} "
             f"array_length: {self.array_length} "
             f"enum_table: {self.enum_table} "
             f"member_offset: {self.member_offset} "
             f"parent: {self.parent.name if self.parent is not None else None} "
             f"attributes: {self.attributes} "
             f"data_type: " + (f"{self.data_type:<32}" if isinstance(self.data_type, str) else ("\n" + str(self.data_type))) +
             f"\n")
        for child in self.children:
            s += str(child)
        return s

    __repr__ = __str__


class Parser:
    def __init__(self, file):
        self.file = file
        self.parse = ELFFile(open(file, 'rb'))
        set_global_machine_arch(self.parse.get_machine_arch())
        self.dbg_info = self.parse.get_dwarf_info()
        self.vars = {}

    def get_header(self):
        print("header", self.parse.header)
        return self.parse.header

    def get_machine_arch(self):
        print(self.parse.get_machine_arch())
        return self.parse.get_machine_arch()

    def read_sections(self):
        """
        Sections描述了链接过程中的需要的符号表、数据、指令等信息。
        """
        ...

    def read_segments(self):
        """
        segments内核使用，内存映射，提供运行时的信息,是经过合并的Sections。
        """
        ...

    def read_symbols(self):
        """
        符号表包含定位和重定位程序的符号定义和引用所需的信息.
        包含变量和函数的符号表信息。
        readelf -s -W Phigent.elf > 456.txt
        """
        ...

    def read_string_tables(self):
        """
        字符串表存储字符串信息的表格。
        sec['sh_name']是个索引，通过sec_string_table.get_string(sec['sh_name']) 可读取section的名称，和已经解析好的sec.name相同。
        """
        ...

    def read_debug_abbrev(self):
        """
        调试信息的缩写表，描述了调试信息中使用的标签和属性。
        """
        ...

    @staticmethod
    def check_cu_has_var(cu):
        """
        查看CU中是否含有变量。
        """
        abbrev_table = cu.get_abbrev_table()
        for code in range(1, 100):  # 可能最大38
            try:
                abbrev = abbrev_table.get_abbrev(code)
                if abbrev.decl.tag == "DW_TAG_variable":
                    return True
            except KeyError:
                return False
        return False

    def has_debug_info(self):
        """
        判断ELF文件是否有调试信息。
        """
        return self.parse.has_dwarf_info()

    def read_variables(self):
        """
        从调试信息中读取变量信息。
        如果ELF文件没有调试信息，则从符号表中读取变量信息。
        """
        if self.vars:
            return self.vars

        if self.has_debug_info():
            self.read_variables_in_debug_info()
        else:
            self.read_variables_in_symtab()
        print(f"\nLoading variables finish.")

    def read_variables_in_debug_info(self):
        """
        从调试信息中读取变量信息。
        调试信息中包含了变量的类型、名称、地址、大小、数组长度、枚举表等信息。
        readelf -w -W Test.elf > debug_info.txt
        """
        len_CU = len(list(self.dbg_info.iter_CUs()))
        for index, cu in enumerate(self.dbg_info.iter_CUs()):
            print(f"\rLoading variables in ELF Compilation Unit {index + 1}/{len_CU}", end="")
            if index < 260:
                continue

            # 跳过没有变量的CU, 节省时间。
            if not self.check_cu_has_var(cu):
                continue

            # for index_die, die in enumerate(cu.iter_DIEs()):  # iter_DIEs()这里会有很多重复的die，不用这个了。
            for index_die, die in enumerate(cu.get_top_DIE().iter_children()):
                if die.tag == "DW_TAG_variable":
                    # DW_AT_location代表变量的地址，是定义变量的地方，声明的地方没有变量地址。
                    if 'DW_AT_location' in die.attributes:
                        name = self.DW_AT_name(die)
                        addr = self.DW_AT_location(die)
                        dt = self.read_dt(self.read_die_value(die, 'DW_AT_type'))
                        self.vars[name] = Item(name=name, addr=addr, data_type=dt)
                        LOG.debug(f"\n{index} {name}: {self.vars[name]}")

    def read_variables_in_symtab(self):
        """
        符号表包含定位和重定位程序的符号定义和引用所需的信息.
        包含变量和函数的符号表信息。
        readelf -s -W Test.elf > symtab.txt
        """
        for sec in self.parse.iter_sections('SHT_SYMTAB'):  # 遍历符号表，一般只有一个。
            num_symbols = sec.num_symbols()
            for index, symbol in enumerate(sec.iter_symbols()):
                print(f"\rLoading variables in ELF Symbol Table {index + 1}/{num_symbols}", end="")
                if (symbol['st_info']['type'] == 'STT_OBJECT'  # 变量
                        and symbol['st_info']['bind'] == 'STB_GLOBAL'):
                    name = symbol.name
                    addr = symbol.entry.st_value
                    data_type = ""  # unknown
                    self.vars[name] = Item(name, addr, data_type)

    @staticmethod
    def read_die_value(die, DW):
        """
        读取DIE中的属性值。
        """
        if DW in die.attributes:
            return die.attributes[DW].raw_value
        else:
            return None

    def DW_AT_name(self, die):
        """
        读取DIE的名称，从byte转为str。
        """
        return self.read_die_value(die, 'DW_AT_name').decode('gbk', errors='ignore')

    def DW_AT_location(self, die):
        """
        读取DIE中变量的地址信息，转成int。
        """
        value = self.read_die_value(die, 'DW_AT_location')
        if not isinstance(value, list):
            LOG.error(f"Error format/length location {value}. {die}")
            return 0
        if len(value) == 5:
            # value[0]不知道啥意思。
            addr = value[1] | (value[2] << 8) | (value[3] << 16) | (value[4] << 24)
            return addr
        else:
            LOG.error(f"Error format/length location {value}. {die}")
            return 0

    def DW_AT_data_member_location(self, die):
        """
        读取DIE中结构体成员的地址信息，转成int。
        """
        value = self.read_die_value(die, 'DW_AT_data_member_location')
        if not isinstance(value, list):
            LOG.error(f"Error format/length location {value}. {die}")
            return 0
        if len(value) == 2:
            return value[1]
        elif len(value) == 3:
            return value[1]
        else:
            LOG.error(f"Error format/length location {value}. {die}")
            return 0

    def read_dt(self, refaddr, parent=None):
        """
        读取数据类型。
        :param refaddr: (int)引用数据类型偏移。
        :param parent: (Item)结构体或数组的父节点。(有对父节点的修改)
        :return: (str/Item)字符串表示基本类型，Item表示结构体数组等复杂类型。
        """
        if parent is None:
            parent = Item()

        die = self.dbg_info.get_DIE_from_refaddr(refaddr)

        if die.tag == "DW_TAG_base_type":
            parent.data_type_size = int(self.read_die_value(die, "DW_AT_byte_size"))
            ret = self.DW_AT_name(die)
        elif die.tag == "DW_TAG_structure_type":
            # 直接修改父节点，member放到children中。
            obj = parent
            obj.data_type_size = self.read_die_value(die, "DW_AT_byte_size")
            if die.has_children:
                for child in die.iter_children():
                    member = self.read_dt(child.offset, obj)
                    member.parent = obj
                    member.level = obj.level + 1
                    obj.append(member)
            else:
                LOG.error("No children in structure_type", die)
            ret = obj
        elif die.tag == "DW_TAG_member":
            # 结构体成员。
            obj = Item()
            obj.name = self.DW_AT_name(die)
            obj.data_type_size = int(self.read_die_value(die, "DW_AT_byte_size"))
            obj.member_offset = self.DW_AT_data_member_location(die)
            obj.level = parent.level + 1
            dt = self.read_dt(die.attributes['DW_AT_type'].raw_value + die.cu.cu_offset, obj)
            # 如果当前成员的数据类型是结构体，则data_type的名称是默认的。否则会出现循环引用，即obj.data_type = obj。
            if isinstance(dt, str):
                obj.data_type = dt
            ret = obj
        elif die.tag == "DW_TAG_array_type":
            # 直接修改父节点，array_length表示数组长度。
            obj = parent
            dt = self.read_dt(die.attributes['DW_AT_type'].raw_value + die.cu.cu_offset, obj)
            if isinstance(dt, str):
                obj.data_type = dt
            obj.array_length = int(int(self.read_die_value(die, "DW_AT_byte_size")) / obj.data_type_size)
            obj.level += 1
            ret = obj
        elif die.tag == "DW_TAG_enumeration_type":
            # 枚举类型。
            obj = parent
            obj.data_type_size = int(self.read_die_value(die, "DW_AT_byte_size"))
            obj.enum_table = {}
            for child in die.iter_children():
                if child.tag == "DW_TAG_enumerator":
                    name = self.DW_AT_name(child)
                    value = int(self.read_die_value(child, "DW_AT_const_value"))
                    obj.enum_table[value] = name
            ret = obj
        elif die.tag == "DW_TAG_typedef":
            ret = self.read_dt(die.attributes['DW_AT_type'].raw_value + die.cu.cu_offset, parent)
        elif die.tag == "DW_TAG_const_type":
            parent.attributes.append("const")
            ret = self.read_dt(die.attributes['DW_AT_type'].raw_value + die.cu.cu_offset, parent)
        elif die.tag == "DW_TAG_volatile_type":
            parent.attributes.append("volatile")
            ret = self.read_dt(die.attributes['DW_AT_type'].raw_value + die.cu.cu_offset, parent)
        elif die.tag == "DW_TAG_pointer_type":
            parent.attributes.append("pointer")
            parent.data_type = "pointer TBD"
            # 链表会出现循环引用，需要跳过
            # ret = self.read_dt(die.attributes['DW_AT_type'].raw_value + die.cu.cu_offset, parent)
            ret = Item()
        elif die.tag == "DW_TAG_subroutine_type":
            # 函数类型, 跳过。
            parent.attributes.append("function")
            ret = Item()
        elif die.tag == "DW_TAG_formal_parameter":
            # 函数类型中的参数, 跳过。
            parent.attributes.append("argument")
            ret = Item()
        else:
            ret = Item()
        return ret
